import cv2 as cv
from matplotlib import pyplot as plt
import numpy as np
import time
from core.utils import write_sequence, lines_sort, relative_position_repair
from PIL import Image
from tqdm import tqdm

Img_dir = "D:/Miniconda3/envs/pyQtApp_00/Data/grbl_Writter/images/test_000.png"
DebugMode = bool(0)

class gcode_generate(): # 计算速度待增加
    def __init__(self, Grid):
        self.move_speed = 1e4
        self.write_speed = 1e4
        self.up_pos = -2.0 # 物理限位到底同时笔也到底 # 刚性笔，Z回原点后笔尖落地拧紧 # 上限位，Z回原点后下降5mm上笔，笔尖落地拧紧
        # self.down_pos = -0.5
        self.down_pos = -6.2
        self.Grid = Grid # 图片渲染缩放系数
        self.zoomcal = 1.0
        self.y_Adjust = 1.05 # 该变量与生成的Y偏移绑定，用于识别是否处于一行
        self.frame = None
        self.dilate = 22 # 膨胀 注意: 最终版中需与生成的文字大小匹配, 达成框选一个字的效果
        self.cv_frame = 3000 # 提供一个合适的缩放比例用于识别文字
        self.search_radius = 15 # 点集连线问题查找半径，越大越快，但会有笔画不全的问题
        self.resultdir = "D:/Miniconda3/envs/pyQtApp_00/Data/grbl_Writter/results"
        self.machine_Adjust_x = 0.0
        self.machine_Adjust_y = 0.0
        self.machine_trans_x = 0.0
        self.machine_trans_y = 0.0
        self.machine_trans_theta = -90.0
        self.machine_grid_compensation = [ # 床面补偿
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [-0.1, 0.0, 0.0]
        ]

    # 简单的先前处理
    def dilate_demo(self) -> np.ndarray:
        print("frame.shape: ",self.frame.shape)
        gray = cv.cvtColor(self.frame, cv.COLOR_BGR2GRAY) # 转换为灰度图
        ret, dst = cv.threshold(gray, 0, 150, cv.THRESH_BINARY_INV | cv.THRESH_OTSU) # 二值化
        num = np.argmax(np.array(self.frame.shape))
        self.zoomcal = self.Grid / (self.frame.shape[num] / self.cv_frame) # 计算图像处理缩放系数
        dst = cv.resize(dst, (0, 0), fx=(self.cv_frame / self.frame.shape[num]), fy=(self.cv_frame / self.frame.shape[num]))
        self.frame = cv.resize(self.frame, (0, 0), fx=(self.cv_frame / self.frame.shape[num]), fy=(self.cv_frame / self.frame.shape[num]))

        # 噪声去除
        morphologyEx_kernel = np.ones((2,2),np.uint8)
        dst = cv.morphologyEx(dst, cv.MORPH_OPEN, morphologyEx_kernel, iterations=1)

        # 侵蚀
        # erode_kernel = np.ones((2, 2), np.uint8)
        # dst = cv.erode(dst, erode_kernel, iterations=1)

        # 边缘平滑
        filter2D_kernel = np.ones((2,2),np.uint8)
        dst = cv.filter2D(dst, -1, filter2D_kernel)
        if (DebugMode):
            plt.imshow(dst)
            plt.show()

        return dst

    # 膨胀然后框选每一个字
    def Contour_demo(self, image) -> list:
        boxes = []
        # 膨胀 注意: 最终版中需与生成的文字大小匹配, 达成框选一个字的效果
        shape = int(image.shape[0] / 5000 * self.dilate)
        kernel = np.ones((shape, shape), np.uint8)
        image = cv.dilate(image, kernel, iterations=1)
        # 需要轮廓层级的轮廓识别
        contours, hierarchy = cv.findContours(image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)

        # 轮廓转方框
        for i in range(0, len(contours)):
            x, y, w, h = cv.boundingRect(contours[i])
            boxes.append([x, y, w, h])

        return boxes

    # 选取较大的框作为字
    def area_cal_choose(self, boxes: list) -> list:
        areas = []
        for box in boxes:
            areas.append(box[2] * box[3])
        # 按框面积排序
        arg = np.argsort(np.array(areas)).tolist()
        arg = arg[::-1]

        # 类似NMS的方法求需要的大框
        average = areas[arg[0]]
        result = []
        for i in arg:
            if (areas[i] / average < 0.3):
                break
            average += areas[i]
            average /= 2
            result.append(boxes[i])

        return result

    # 对框进行排序
    def area_sort(self, boxes: list) -> list:
        boxes = np.array(boxes)
        boxes = boxes[boxes[:, 1].argsort()] # 按照y进行排序
        # boxes = boxes[boxes[:, 0].argsort()]  # 按照x进行排序

        x_sort = []
        counter = 0
        while (counter <= len(boxes) - 1):   # 遍历boxes
            average = boxes[counter][1]
            y_sort = []
            while (counter <= len(boxes) - 1):   # 判断同行box
                x, y, w, h = boxes[counter]
                if (y / average > self.y_Adjust):        # 该变量与生成的Y偏移绑定
                    break
                average += y
                average /= 2
                y_sort.append([x, y, w, h])
                counter += 1
            y_sort = np.array(y_sort)   # 同行根据x排序
            y_sort = y_sort[y_sort[:, 0].argsort()]
            y_sort = y_sort.tolist()
            x_sort += y_sort

        print("Num of contours: ", len(x_sort))
        return x_sort

    # 细化求点集
    def thin_demo(self, image) -> list:
        # 细化
        thinned = cv.ximgproc.thinning(image, thinningType=cv.ximgproc.THINNING_ZHANGSUEN)
        # contours, hierarchy = cv.findContours(thinned, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)

        # 选取
        thinned = np.argwhere(thinned > 0)
        thinned[:, [0, 1]] = thinned[:, [1, 0]]

        if (DebugMode):
            show = np.zeros(image.shape[:2], np.uint8)
            # 打点
            for point in thinned:
                cv.circle(show, point, 1, (255, 255, 255))
                # cv.arrowedLine(show, (0, 0), (80, 80), (255, 255, 255), 1, 0, 0, 0.07)
            plt.imshow(show)
            plt.show()
        # print("Num of dots: ", len(thinned))
        return thinned.tolist()

    def Debug_show_lines(self, image, lines):
        if (DebugMode):
            show = np.zeros(image.shape[:2], np.uint8)
            # 打点
            for line in lines:
                for point in line:
                    cv.circle(show, point, 1, (255, 255, 255))
                    # cv.arrowedLine(show, (0, 0), (80, 80), (255, 255, 255), 1, 0, 0, 0.07)
            plt.imshow(show)
            plt.show()

    def Rotate_90_degrees(self, lines: list) -> list:
        lines_backup = []
        for line in lines:
            line_backup = []
            for point in line: # 对每一个点左乘旋转矩阵并平移frame边长
                point_backup = np.dot(self.matrix, np.array(point)) + np.array([0, self.frame.shape[1]])
                point = point_backup.tolist()
                line_backup.append(point)
            lines_backup.append(line_backup)
        return lines_backup

    def z_grid_compensation(self, point): # 二分法求当前点的分区 # ...默认为逆时针旋转90度
        grid = self.frame.shape[0] / len(self.machine_grid_compensation)
        count0 = len(self.machine_grid_compensation) // 2
        cloest_grid = grid * count0
        while (abs(point[0] - cloest_grid) > grid):
            if ((point[0] > cloest_grid)):
                count0 += count0 // 2 + 1
            elif ((point[0] < cloest_grid)):
                count0 -= count0 // 2
            cloest_grid = grid * count0
        if (point[0] > cloest_grid):
            count0 += 1

        grid = self.frame.shape[1] / len(self.machine_grid_compensation[0])
        count1 = len(self.machine_grid_compensation[0]) // 2
        cloest_grid = grid * count1
        while (abs(point[1] - cloest_grid) > grid):
            if ((point[1] > cloest_grid)):
                count1 += count1 // 2 + 1
            elif ((point[1] < cloest_grid)):
                count1 -= count1 // 2
            cloest_grid = grid * count1
        if (point[1] > cloest_grid):
            count1 += 1

        return [count0, count1]

    # 线段转Gcode
    def result1(self, vactors, filename: str):
        with open(filename, "w") as f:
            # f.write("$H\n")   # 启动
            f.write("G21\n")
            f.write("G90\n")
            f.write(f"G1 G90 Z{self.up_pos} F{self.move_speed}\n")   # 抬笔
            for vactor in vactors:
                h = self.z_grid_compensation(vactor[0])
                f.write("G0 X%5.3f F%d\n" % (vactor[0][0] / self.zoomcal + self.machine_Adjust_x, \
                                                    self.move_speed))  # 移动到目标位置
                f.write("G0 Y%5.3f F%d\n" % (-vactor[0][1] / self.zoomcal + self.machine_Adjust_y, \
                                                    self.move_speed))  # 移动到目标位置
                f.write(f"G1 G90 Z%5.3f F{self.move_speed}\n" % (self.down_pos + \
                    self.machine_grid_compensation[h[0]-1][len(self.machine_grid_compensation[0])-h[1]]))  # 落笔
                for vac in vactor:
                    f.write("G1 X%5.3f Y%5.3f F%d\n" % (vac[0] / self.zoomcal + self.machine_Adjust_x, \
                                                        -vac[1] / self.zoomcal + self.machine_Adjust_y, \
                                                        self.write_speed))  # 书写
                f.write(f"G1 G90 Z{self.up_pos} F{self.move_speed}\n")   # 抬笔
            f.write("G90 G0 X10.0 Y-10.0\n")
        f.close()

    def import_image_from_file(self, Img_dir):
        self.frame = cv.imread(Img_dir)  # 读取图片

    def import_image_from_Image(self, image:Image.Image): # 格式转换并导入
        self.frame = cv.cvtColor(np.array(image), cv.COLOR_RGB2BGR)

    def gcode_generate_and_result(self, filename: str):
        vactors = []
        self.matrix = np.array([[0, 1], [-1, 0]])  # 逆时针旋转90度矩阵
        if (DebugMode):
            self.contours = self.contours[9:]
        pbar = tqdm(self.contours)
        for i, contour in enumerate(pbar):  # 将每一个框中的文字分解为 点、
            x, y, w, h = contour
            word = self.erodeed_img[y:(y + h), x:(x + w)]
            lines = []
            dots = self.thin_demo(word)  # 细化求点集
            if (len(dots) > 0):
                lines = write_sequence(dots, self.search_radius)  # 点集连线问题主函数
                if (len(lines) > 0):
                    lines = lines_sort(lines)  # 线段排序
                    self.Debug_show_lines(word, lines) # 用于debug点集连线
                    lines = relative_position_repair(lines, contour)  # 点集转换为绝对坐标
                    lines = self.Rotate_90_degrees(lines)  # 逆时针旋转90度

            if (DebugMode):
                plt.imshow(word)
                plt.show()

            pbar.set_description(f"dots: {len(dots)}, vactors: {len(lines)}")
            vactors += lines

        self.result1(vactors, filename)  # 线段转Gcode

    def Img2Gcode_main(self): # Img2Gcode主函数返回识别后图片

        self.erodeed_img = self.dilate_demo() # 简单的先前处理

        contours = self.Contour_demo(self.erodeed_img) # 膨胀然后框选每一个字

        contours = self.area_cal_choose(contours) # 选取较大的框作为字

        contours = self.area_sort(contours) # 对框进行排序

        self.contours = contours

        for i, contour in enumerate(contours):
            x, y, w, h = contour
            cv.putText(self.frame, f'{i}', (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.8, (153, 153, 0), 2, cv.LINE_AA)
            cv.rectangle(self.frame, (x, y), (x + w, y + h), (153, 153, 0), 2)

        self.frame = cv.cvtColor(self.frame, cv.COLOR_BGR2RGB)

        return self.frame

if __name__ == '__main__':
    run = gcode_generate(30.0)
    run.import_image_from_file("D:/Miniconda3/envs/pyQtApp_00/Data/grbl_Writter/images/test_001_0_1_1.png")
    run.Img2Gcode_main()
    run.gcode_generate_and_result(run.resultdir + f"\\result{time.time()}.nc")
